from typing import Optional
from typing import Dict

from aioredis import from_url
from aioredis import Redis
from loguru import logger

from common.config import settings


class RedisClient:

    def __init__(self):
        self._pool: Dict[int, Redis] = {
            db: from_url(settings.REDIS_URL, db=db, max_connections=10, encoding="utf-8",
                         decode_responses=True)
            for db in range(16)
            }

    def use(self, db: int = 0) -> Optional[Redis]:
        """切换redis db"""
        logger.debug(f'切换Redis DB {db}')
        return self._pool.get(db)

    async def close(self):
        for v in self._pool.values():
            await v.close()


def use_redis(request, db: int) -> Redis:
    """从fastapi 实例中切换redis DB
    :param request: FastAPI Request对象
    :param db: DB
    :return Redis 对象
    """
    return request.app.state.redis_cli.use(db)
